from utils.cloud_utils import notify_cloud_v2
from config import similarity_threshold
from celery import chain, group
from celerytask import app
from celery import chord
from utils.video_utils import synthesize_video
from tybase.tools.image_utils import delete_file_or_directory
from tybase.tools.image_utils import url_to_md5
from utils.cloud_utils import upload_file_to_aliyun, transform_path, upload_file_to_qiniu
from loguru import logger
from utils.video_utils import split_video_into_frames, 加水印
from config import kLocUserImagesPath, kSourceVideosFrame, kUserSwapperFaceFrames, kFontPath
import time
import os
import sys
sys.path.append("..")  # 也可以这样
import config
# 人脸验证的api流程
preprocess_source_images_v2 = app.tasks['celerytask.tasks.preprocess_source_images_v2']
process_workflow_V3 = app.tasks['celerytask.tasks.start_V3']  # 单人换脸的celery逻辑
# 把这整个包装成一个task,然后直接发送任务就可以了


@app.task(soft_time_limit=60, time_limit=120)
def async_split_video_into_frames(video_url, frame_rate, frame_name, save_path,task_id):
    # 分割视频
    try:
        video_loc_frame_path = split_video_into_frames(video_url,
                                                       kSourceVideosFrame=save_path,
                                                       frames=frame_rate,
                                                       frame_name=frame_name)

        # 准备异步任务数据
        full_video_paths_list = [os.path.join(video_loc_frame_path, i) for i in os.listdir(
            video_loc_frame_path) if ".aac" not in i and ".jpg" in i]
        output_directory_name = f"{video_loc_frame_path}_{int(time.time())}"
        # 返回需要的数据
        return {
            "video_loc_frame_path": video_loc_frame_path,
            "full_video_paths_list": full_video_paths_list,
            "output_directory_name": output_directory_name
        }
    except Exception as e:
        # 出现异常,先删除掉产生的文件夹
        delete_file_or_directory(video_loc_frame_path)
        # delete_file_or_directory(video_url)
        # 这是需要传递的参数
        notify_cloud_v2(result={},
                        task_id=task_id,
                        status=250,
                        error_message="视频分割失败"
                        )
        raise e  # 重新抛出异常以通知 Celery 任务失败
    finally:
        pass
        # delete_file_or_directory(video_url)


# 换脸
@app.task(bind=True)
def process_images_many_face_task(self, face_pairs_list, full_video_paths, output_directory_name, GFPGANAmount,task_id):
    # 多张图片换脸
    # 多人视频换脸的核心逻辑,这里看下,修改一下就可以做成单图换脸的逻辑了!
    try:
        os.makedirs(output_directory_name, exist_ok=True)
        signatures = []

        for target_file in full_video_paths:
            output_file_path = os.path.join(
                output_directory_name, os.path.basename(target_file))
            signatures.append(process_workflow_V3.s(
                target_file, face_pairs_list, output_file_path, GFPGANAmount, similarity_threshold))

        # 创建一个group，里面有你的所有任务签名
        job = group(signatures)

        # 执行这个group，并异步等待结果
        result = job.apply_async()

        # 这里我们不在任务内部等待结果，而是返回AsyncResult的id
        return result.id
    except Exception as e:
        # 出现异常,先删除掉产生的文件夹
        delete_file_or_directory(output_directory_name)
        # 这是需要传递的参数
        notify_cloud_v2(result={},
                        task_id=task_id,
                        status=250,
                        error_message="换脸过程失败"
                        )
        logger.error(e)
        raise e

@app.task(bind=True,soft_time_limit=60, time_limit=120)
def post_process_task(self, _ignored_result, output_directory_name, video_loc_frame_path, requests_data, 帧数, frame_name, 水印, oss_name,task_id):
    # 最终的组合发布任务
    # 换脸完成以后需要做的任务
    try:
        output_path = kUserSwapperFaceFrames + \
            f"/{url_to_md5(requests_data['target_video']).split('.')[0]}_{int(time.time())}.mp4"
        logger.info("换脸完毕,合成视频路径:{}".format(output_path))
        synthesize_video(image_directory=output_directory_name,
                         audio_path=video_loc_frame_path + "/audio.aac",
                         output_path=output_path,
                         frame_rate=帧数,
                         crf=config.crf,
                         frame_name=frame_name,
                         )

        logger.info("换脸完毕,合成视频路径:{}".format(output_path))
        print("要不要加水印?")
        if 水印:
            加水印(output_path, fontfile_path=kFontPath)

        aliyun_file_path = transform_path(output_path)

        # # 把链接上传到阿里云,构建好完整的路径

        if oss_name == "qiniu":
            cloud_url = res_url = upload_file_to_qiniu(
                aliyun_file_path, output_path)
        else:
            cloud_url = upload_file_to_aliyun(
                aliyun_file_path, output_path)
            res_url = cloud_url.split('aliyuncs.com/')[-1]
            # 这里执行你的实际工作
        res = {
            "video_loc_frame_path": video_loc_frame_path,
            "output_directory_name": output_directory_name,
            "output_path": output_path,
            "aliyun_file_path": cloud_url,
            "ty_path": res_url,
            "verify_face": True,
            "error_message": ""
        }
        notify_cloud_v2(result=res,
                        task_id=task_id,
                        status=100,
                        error_message=""
                        )
    except Exception as e:
        notify_cloud_v2(result={},
                        task_id=task_id,
                        status=250,
                        error_message="视频合成失败"
                        )
        raise e

    finally:
        delete_file_or_directory(output_path)
        delete_file_or_directory(output_directory_name)
        delete_file_or_directory(video_loc_frame_path)
    return res


# 创建一个回调函数，它将处理 task1 和 task2 的结果
@app.task
def callback(results, requests_data, 帧数, frame_name, 水印, oss_name,task_id):
    # 组装任务
    result_task1, result_task2 = results
    item = {**result_task1, "face_pairs_list": result_task2}
    followup_chain = chain(
        process_images_many_face_task.s(item["face_pairs_list"],
                                        item["full_video_paths_list"], item["output_directory_name"], GFPGANAmount=100,task_id=task_id),
        # 视频合成的部分
        post_process_task.s(item["output_directory_name"], item["video_loc_frame_path"],
                            requests_data, 帧数, frame_name, 水印, oss_name,task_id)

    )
    return followup_chain.apply_async()
